import asyncio
import json

# 用来发起异步HTTP请求
import aiohttp


class TraditionalMCPClient:
    def __init__(self, server_url: str = "http://127.0.0.1:8000"):
        self.server_url = server_url
        self.session_id = None  # 自己的MCP会话ID
        self.sse_task = None  # SSE的监听任务
        self.session = None  # aiohttp会话对象

    async def __aenter__(self):
        self.session = aiohttp.ClientSession()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self.sse_task:
            self.sse_task.cancel()  # 取消SSE监听任务
        if self.session:
            await self.session.close()  # 关闭 aiohttp会话

    async def shutdown(self) -> bool:
        if not self.session_id:
            return True
        try:
            payload = {"jsonrpc": "2.0", "id": 4, "method": "shutdown", "params": {}}
            headers = {"X-Session-ID": self.session_id}
            async with self.session.post(
                f"{self.server_url}/shutdown", json=payload, headers=headers
            ) as response:
                if response.status == 200:
                    print("连接关闭成功")
                    self.session_id = None
                    return True
                else:
                    print(f"连接关闭失败:{response.status}")
                    return False

        except Exception as e:
            print(f"连接关闭异常:{e}")
            return False

    async def initialize(self) -> bool:
        try:
            payload = {
                "jsonrpc": "2.0",
                "id": 1,
                "method": "initialize",
                "params": {
                    "protocol_version": "2024-11-05",
                    "capabilities": {"tools": {}, "resources": {}, "prompts": {}},
                    "client_info": {
                        "name": "Traditional MCP Client",
                        "version": "1.0.0",
                    },
                },
            }
            async with self.session.post(
                f"{self.server_url}/initialize", json=payload
            ) as response:
                if response.status == 200:
                    result = await response.json()
                    self.session_id = result["result"]["serverInfo"]["session_id"]
                    print("连接初始化成功")
                    return True
                else:
                    print(f"连接初始化的失败:{response.status}")
                    return False

        except Exception as e:
            print(f"连接 初始化的异常:{e}")
            return False

    async def establish_sse_connection(self) -> bool:
        if not self.session_id:
            print(f"请先初始化连接")
            return False
        try:
            headers = {"X-Session-ID": self.session_id}
            async with self.session.get(
                f"{self.server_url}/sse", headers=headers
            ) as response:
                if response.status == 200:
                    print("SSE连接成功")
                    # 启动SSE的监听任务
                    self.sse_task = asyncio.create_task(self._listen_sse(response))
                    return True
                else:
                    print(f"SSE建立连接错误:{response.status}")
                    return False
        except Exception as e:
            print(f"SSE建立连接异常:{e}")
            return False

    async def _listen_sse(self, response):
        try:
            async for line in response.content:
                line = line.decode("utf-8").strip()
                if line.startswith("data: "):
                    try:
                        data = json.loads(line[6:])
                        if "method" in data:
                            print(f"SSE消息:{data['method']}")
                    except json.JSONDecodeError:
                        continue
        except asyncio.CancelledError:
            print(f"SSE连接已经取消")
        except Exception as e:
            if "Connection closed" in str(e) or "Connection reset" in str(e):
                print("SSE流正常关闭")
            else:
                print(f"SSE监听异常:{str(e)}")

    async def _handle_sse_response(self, response):
        result = None
        try:
            async for line in response.content:
                line = line.decode("utf-8").strip()
                if line.startswith("data: "):
                    try:
                        data = json.loads(line[6:])
                        if "method" in data:
                            if "notification/progress" in data["method"]:
                                progress_data = data["params"]
                                if progress_data["type"] == "begin":
                                    print(f"任务开始:{progress_data["title"]}")
                                elif progress_data["type"] == "update":
                                    print(f"进度更新:{progress_data["message"]}")
                                elif progress_data["type"] == "end":
                                    print(f"任务完成:{progress_data["title"]}")
                            elif "notifications/heartbeat" in data["method"]:
                                print(f"心跳:{data['params'].get('timestamp','N/A')}")
                        elif "result" in data:
                            result = data["result"]["content"][0]["text"]
                            break
                    except json.JSONDecodeError as e:
                        print(f"SSE数据解析失败:{str(e)},原始数据{line}")
                        continue
        except Exception as e:
            print(f"处理SSE流异常:{str(e)}")
        return result or {"error": "No result received from SSE stream"}

    async def call_tool(self, tool_name: str, arguments):
        if not self.session_id:
            print("请先初始化连接")
            return None
        try:
            payload = {
                "jsonrpc": "2.0",
                "id": 3,
                "method": "tools/call",
                "params": {"name": tool_name, "arguments": arguments},
            }
            headers = {"X-Session-ID": self.session_id}
            async with self.session.post(
                f"{self.server_url}/tools/call", json=payload, headers=headers
            ) as response:
                if response.status == 200:
                    content_type = response.headers.get("content-type", "")
                    if "text/event-stream" in content_type:
                        print(f"工具{tool_name}返回了流式响应")
                        return await self._handle_sse_response(response)
                    else:
                        try:
                            result = await response.json()
                            print(f"工具{tool_name}调用成功")
                            return result["result"]["content"][0]["text"]
                        except Exception as json_error:
                            print(f"解析失败:{json_error}")
                            text_content = await response.text()
                            return {
                                "error": "JSON解析失败",
                                "raw_content": text_content,
                            }

                else:
                    print(f"获取工具列表失败:{response.status}")
                    return []
        except Exception as e:
            print(f"获取工具列表异常:{e}")
            return []

    async def list_tools(self) -> list:
        if not self.session_id:
            print("请先初始化连接")
            return []
        try:
            payload = {"jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {}}
            headers = {"X-Session-ID": self.session_id}
            async with self.session.post(
                f"{self.server_url}/tools/list", json=payload, headers=headers
            ) as response:
                print(response)
                if response.status == 200:
                    result = await response.json()
                    tools = result["result"]["tools"]
                    print(f"获取工具列表成功,获取到了{len(tools)}个工具")
                    for tool in tools:
                        print(f"{tool['name']}:{tool['description']}")
                    return tools
                else:
                    print(f"获取工具列表失败:{response.status}")
                    return []
        except Exception as e:
            print(f"获取工具列表异常:{e}")
            return []


async def traditional_mcp():
    async with TraditionalMCPClient() as client:
        # 1.初始化连接
        if not await client.initialize():
            print("初始化失败")
            return
        # 2.获取工具列表
        tools = await client.list_tools()
        if not tools:
            print("无法获取工具列表")
            return
        # 3.调用工具
        result = await client.call_tool("echo", {"message": "hello"})
        if result:
            print(f"结果:{result}")
        # 4.建立SSE连接
        if await client.establish_sse_connection():
            await asyncio.sleep(1)
        # 5.调用长运行工具
        result = await client.call_tool("long_running_task", {"duration": 5})
        if result:
            print(f"最终结果为:{result}")
        # 6. 关闭连接
        await client.shutdown()


if __name__ == "__main__":
    asyncio.run(traditional_mcp())
